#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import binascii
import hashlib
import os
import sys
import time
import json
import logging
import logging.handlers
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List
import requests
from requests.auth import AuthBase
from tqdm import tqdm 

from utils.api_connector import ApiConnector, ModelConnector

# import utils


def get_logger():
    fmt = '%(asctime)s %(levelname)s[%(filename)s:%(lineno)d][%(funcName)s] %(message)s'
    fmt = logging.Formatter(fmt, '%Y-%m-%d %H:%M:%S')

    folder = os.path.join(os.getcwd(), 'logs')
    if not os.path.exists(folder):
        os.mkdir(folder)
    path = os.path.join(folder, 'iwiki.log')

    file_handler = logging.handlers.TimedRotatingFileHandler(
        path, when='D', backupCount=90, encoding='utf-8'
    )
    file_handler.setFormatter(fmt)
    logger = logging.getLogger('root')

    stream_handler = logging.StreamHandler(sys.stderr)
    stream_handler.setFormatter(fmt)
    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)
    logger.setLevel('DEBUG')
    logger.propagate = False
    return logger


def hexlify(s):
    return binascii.hexlify(s).decode('utf-8')


def calc_signature(token, timestamp, nonce):
    s = u'{timestamp}{token}{nonce}{timestamp}'.format(timestamp=timestamp, token=token, nonce=nonce)
    s = hashlib.sha256(s.encode('utf-8')).digest()
    return hexlify(s)


def prepare_rio_headers(paasid, token):
    timestamp = str(int(time.time()))
    nonce = hexlify(os.urandom(16))
    signature = calc_signature(token, timestamp, nonce)
    return {
        'x-rio-paasid': paasid,
        'x-rio-signature': signature,
        'x-rio-timestamp': timestamp,
        'x-rio-nonce': nonce,
    }


class RioHeaderAuth(AuthBase):
    def __init__(self, paasid, token):
        self.paasid = paasid
        self.token = token

    def __call__(self, r):
        headers = prepare_rio_headers(self.paasid, self.token)
        r.headers.update(headers)
        return r


def dump(rsp):
    print()
    print('status:', rsp.status_code)
    print('headers:')
    for k, v in rsp.headers.items():
        print(k, v)
    print()
    print('body:')
    print(rsp.text)


class ApiSession(requests.Session):
    def __init__(self, base=None, *args, **kwargs):
        super(ApiSession, self).__init__(*args, **kwargs)
        self._base = base.rstrip('/')

    def request(self, method, url, *args, **kwargs):
        if url.startswith('/'):
            url = self._base + url
        return super(ApiSession, self).request(method, url, *args, **kwargs)


class Wiki(object):
    @classmethod
    def from_env(cls):
        kwargs = {
            'env': os.getenv('env', 'devnet'),
            'paasid': os.getenv('paasid', ''),
            'token': os.getenv('token', ''),
        }
        return Wiki(**kwargs)

    @classmethod
    def from_file(cls, path):
        with open(path, 'r') as fp:
            data = fp.read()
        kwargs = {}
        if path.endswith('.json'):
            kwargs = json.loads(data)
        else:
            for line in data.splitlines(keepends=False):
                args = line.split('=', maxsplit=2)
                if line.startswith('#') or len(args) != 2:
                    continue
                kwargs[args[0]] = args[1]
        kwargs.setdefault('env', 'devnet')
        return Wiki(**kwargs)

    def __init__(self, env, paasid, token):
        regions = {
            'oa':     'http://api-g.sgw.woa.com',
            'oss':    'http://api-idc.sgw.woa.com',
            'idc':    'http://api-g-idc.sgw.woa.com',
            'idcw':   'http://api-g-idcw.sgw.woa.com',
            'devnet': 'http://api-idc.sgw.woa.com',
            'test':   'http://test.devnet.rio.tencent.com',
        }
        base = regions.get(env, None)
        if not base:
            regions = "/".join(regions.keys())
            raise ValueError(f'invalid env "{env}", must be one of {regions}, suggest you try env devnet/idc')
        self.base = f'{regions[env]}/ebus/iwiki/prod'
        s = ApiSession(self.base)
        # s.trust_env = False
        s.auth = RioHeaderAuth(paasid, token)
        self.client = s

    def page_link(self, id, title=None):
        url = f'https://iwiki.woa.com/p/{id}'
        return f'[{title}]({url})' if title else url

    def get_json(self, api, params=None):
        rsp = self.client.get(api, params=params)
        return rsp.json()

    def post_json(self, api, body, params=None):
        rsp = self.client.post(api, json=body, params=params)
        return rsp.json()

    def ping(self):
        return self.client.get('/tencent/api/v2/user/current')

    def get_user_info(self):
        return self.get_json('/tencent/api/v2/user/current')['data']

    def get_doc_body(self, id):
        return self.get_json(f'/tencent/api/v2/doc/body?id={id}&bodymode=md')
    
    def get_doc_body_md(self, id):
        """
        获取文档内容，返回Markdown格式的文本
        """
        data = self.get_json(f'/tencent/api/v2/doc/body?id={id}&bodymode=md')
        if 'data' in data and 'body' in data['data']:
            return data['data']['body']
        else:
            raise ValueError(f'获取文档内容失败: {data.get("msg", "未知错误")}')
        
    
    def get_doc_body_html(self, id):
        """
        获取文档内容，返回HTML格式的文本
        """
        data = self.get_json(f'/tencent/api/v2/doc/body?id={id}&bodymode=view')
        if 'data' in data and 'body' in data['data']:
            return data['data']['body']
        else:
            raise ValueError(f'获取文档内容失败: {data.get("msg", "未知错误")}')

    def get_space(self, k):
        try:
            data = self.get_json(f'/tencent/api/space/{k}')
            return data['result']
        except Exception as e:
            msg = f'get space "{k}" failed: {e}'
            logger.exception(msg)
            raise Exception(msg)

    def children(self, id):
        try:
            data = self.get_json(f'/tencent/api/v2/pagetree/children?parentid={id}')
            return data['data']['list']
        except Exception as e:
            msg = f'get children of page "{id}" failed: {e}'
            logger.exception(msg)
            raise Exception(msg)

    def children_shorter(self, id):
        return [
            {'id': i['id'], 'title': i['title']}
            for i in self.children(id)
        ]

    def find_children(self, id, title):
        children = self.children(id)
        return [i['id'] for i in children if i['title'] == title]

    def mkdir(self, parent, title):
        body = {
            'parentid': parent,
            'name': title,
        }
        try:
            data = self.post_json('/tencent/api/v2/doc/folder/create', body)
            return data['data']['id']
        except Exception as e:
            msg = f'mkdir {title} under {parent} failed: {e}'
            logger.exception(msg)
            raise Exception(msg)

    def create_doc(self, parent_id, title, body, content_type):
        body = {
            'parentid': parent_id,
            'contenttype': content_type,
            'title': title,
            'body': body,
        }
        try:
            data = self.post_json('/tencent/api/v2/doc/create', body)
            return data['data']['id']
        except Exception as e:
            msg = f'create_doc {parent_id} / {title} failed: {e}'
            logger.exception(msg)
            raise Exception(msg)

    def update_doc(self, id, title, body):
        body = {
            'id': id,
            'title': title,
            'body': body,
            'version': 1,
            'force': True,
        }
        try:
            data = self.post_json('/tencent/api/v2/doc/save', body)
            if data['msg'] != 'ok':
                raise Exception('update_doc failed: ' + data['msg'])
        except Exception as e:
            msg = f'update_doc "{title}"({id}) failed: {e}'
            logger.exception(msg)
            raise Exception(msg)

    def upload_attachment(self, doc_id, path, name=None):
        if not name:
            name = os.path.basename(path)
        api = f'/tencent/api/attachments/s3/presign?docid={doc_id}'
        rsp = self.get_json(api)
        attachmentid = rsp['data']['attachmentid']
        url = rsp['data']['url']
        rsp = self.client.put(url, data=open(path, 'rb'))
        ok = rsp.status_code == 200
        body = {
            'attachmentid': attachmentid,
            'status': ok and 'ok' or 'failure',
            'filepath': name
        }
        data = self.post_json('/tencent/api/attachments/notify', body)
        if not ok:
            raise ValueError(f'upload failed: {rsp.text}')
        if data['msg'] != 'ok':
            raise ValueError(f'upload notify failed: {data["msg"]}')
        return attachmentid

class WikiTreeNode:
    def __init__(self, wiki: Wiki, id=4015637120, title="root"):
        self.id = id
        self.title = title
        self.children: List[WikiTreeNode] = []
        self.wiki = wiki
        self.summary = None  # 用于存储文档摘要
        self._lock = threading.Lock()  # 添加线程锁
        
    def add_child(self, child):
        with self._lock:  # 使用锁保护 children 列表
            self.children.append(child)
            
    def get_doc_body(self, style: str = "markdown") -> str:
        """
        获取当前节点文档的内容
        """
        try:
            if style == "markdown":
                return self.wiki.get_doc_body_md(self.id)
            elif style == "html":
                return self.wiki.get_doc_body_html(self.id)
        except Exception as e:
            print(f"获取文档内容失败: {e}")
            return ""

    def count_total_nodes(self, id=None):
        """
        统计以当前节点为根的子树中的总节点数
        """
        if id is None:
            id = self.id
        
        children_data = self.wiki.children_shorter(id)
        total = 1  # 当前节点
        
        for child_data in children_data:
            total += self.count_total_nodes(child_data['id'])
        
        return total

    def construct_subtree(self, api_connector: ModelConnector, pbar=None, verbose=False):
        """
        递归构建子树（多线程版本，带进度条）
        """
        children_data = self.wiki.children_shorter(self.id)
        
        # 多线程并行构建子节点
        def build_child_node(child_data):
            node = WikiTreeNode(child_data['id'], child_data['title'], self.wiki)
            if verbose:
                print(f"Constructing subtree for {node.title} (ID: {node.id})")
            node.construct_subtree(api_connector, pbar, verbose)
            return node
        
        # 使用线程池并行构建所有子节点
        with ThreadPoolExecutor(max_workers=5) as executor:
            # 提交所有子节点构建任务
            future_to_child = {
                executor.submit(build_child_node, child_data): child_data 
                for child_data in children_data
            }
            
            # 等待所有任务完成并收集结果
            for future in as_completed(future_to_child):
                try:
                    node = future.result()
                    self.add_child(node)
                except Exception as e:
                    child_data = future_to_child[future]
                    if verbose:
                        print(f"Error constructing subtree for {child_data['title']}: {e}")
        
        # 所有子节点构建完成后，再执行后续代码
        doc_md = self.wiki.get_doc_body_md(self.id)
        
        if len(self.children) > 0:
            # collect summaries from children
            summaries = []
            for child in self.children:
                if child.summary:
                    summaries.append(child.summary)
            doc_md += "\n\n\n\n".join(summaries)
        
        prompt = ("请对提供的文档进行摘要总结，按照以下模板输出：\n"
                 "---\n"
                 "## 文档标题: {title}\n\n"
                 "## 文档摘要\n\n"
                 "### 1. 主要功能模块\n"
                 "- [功能模块1]\n"
                 "- [功能模块2]\n"
                 "- [功能模块3]\n"
                 "...\n\n"
                 "### 2. 命令接口清单及功能说明\n\n"
                 "| 序号 | 命令接口 | 功能描述 |\n"
                 "|------|----------|----------|\n"
                 "| 1 | `[命令1]` | [功能说明1] |\n"
                 "| 2 | `[命令2]` | [功能说明2] |\n"
                 "| 3 | `[命令3]` | [功能说明3] |\n"
                 "| ... | ... | ... |\n\n"
                 "### 3. 核心特性\n"
                 "- [特性1]\n"
                 "- [特性2]\n"
                 "- [特性3]\n\n"
                 "请确保：\n"
                 "1. 准确提取所有命令行接口\n"
                 "2. 简洁描述每个接口的核心功能\n"
                 "3. 按功能模块分类整理\n"
                 "4. 使用表格形式清晰展示命令与功能的对应关系\n\n"
                 "---\n\n"
                 "{doc_md}").format(doc_md=doc_md, title=self.title)
        
        if verbose:
            print("prompt:", prompt)
        # 调用API获取摘要        
        response = api_connector.send_request(prompt, verbose=verbose)
        self.summary = response
        
        # 更新进度条
        if pbar:
            pbar.update(1)
            pbar.set_description(f"已完成: {self.title}")

    def construct_subtree_with_progress(self, api_connector: ModelConnector, verbose=False):
        """
        带进度条的构建子树方法
        """
        # 首先统计总节点数
        if verbose:
            print("正在统计节点总数...")
        total_nodes = self.count_total_nodes()
        if verbose:
            print(f"发现 {total_nodes} 个节点")
        
        # 创建进度条
        with tqdm(total=total_nodes, desc="构建索引树", unit="节点") as pbar:
            self.construct_subtree(api_connector, pbar, verbose)

    def subtree_to_dict(self):
        """
        将子树转换为字典格式
        """
        return {
            'id': self.id,
            'title': self.title,
            'summary': self.summary,
            'children': [child.subtree_to_dict() for child in self.children]
        }
    
    def from_json_local(self, json_data):
        """
        从本地JSON数据加载子树
        """
        self.id = json_data['id']
        self.title = json_data['title']
        self.summary = json_data.get('summary', None)
        
        for child_data in json_data.get('children', []):
            child_node = WikiTreeNode(self.wiki)
            child_node.from_json_local(child_data)
            self.add_child(child_node)
    
    def from_file(file_path, wiki=None):
        """
        从文件中加载子树
        """
        with open(file_path, 'r') as f:
            data = json.load(f)
            
        local_root = WikiTreeNode(wiki=wiki)
        local_root.from_json_local(data)
        return local_root
        
    def traverse_iteratively(self):
        """
        迭代遍历子树
        :param pbar: 进度条, 默认为None
        """
        stack = [self]
        while stack:
            node = stack.pop()
            yield node
            stack.extend(node.children[::-1])

def gen_index(wiki, root_id):
    """
    生成目录索引
    """
    children = wiki.children_shorter(root_id)
    tree_dict = {}
    
    


if __name__ == '__main__':
    '''
    wiki = Wiki('devnet', 'your-paasid', 'your-token')
    or
    wiki = Wiki.from_env('rio.env')
    # add rio.env file:
    paasid=your-paasid
    token=your-token
    env=devnet
    '''
    logger = get_logger()
    wiki = Wiki.from_file('rio.env')


    # # 验证连接
    # print(wiki.ping().text)
    
    doc_id = 4015637120
    children = wiki.children(doc_id)
    print(json.dumps(children, indent=4, ensure_ascii=False))
    print(children[3]['id'])
    sub_children = wiki.children(children[3]['id'])
    print(json.dumps(sub_children, indent=4, ensure_ascii=False))
    subsub_children = wiki.children(sub_children[0]['id'])
    print(json.dumps(subsub_children, indent=4, ensure_ascii=False))
    print(wiki.get_doc_body(sub_children[0]['id'])['data']['body'])
    
    api_connector = ApiConnector()
    
    treeroot = WikiTreeNode(doc_id, 'root', wiki)
    treeroot.construct_subtree_with_progress(api_connector, verbose=False)
    tree_dict = treeroot.subtree_to_dict()
    with open('tree.json', 'w') as f:
        json.dump(tree_dict, f, indent=4, ensure_ascii=False)





